#! -*- coding: utf-8 -*-
import os
import json
import time
import pandas as pd
from docx import Document

TITLE_LIST = ["资产分布", "借款人分布"]

def load_docx2json(input_file_path):
    doc = Document(input_file_path)

    result_dict = dict()
    final_dict = dict()

    for title in TITLE_LIST:
        result_dict[title] = list()

    for table in doc.tables:
        _title = None
        group_list = list()
        for row in table.rows:
            cell_list = list()
            _index = 1
            for cell in row.cells:
                cell_text = cell.text
                if not cell_text:
                    break
                if cell_text in TITLE_LIST:
                    _title = cell_text
                    break
                if cell_text == "占比" and _index == 1:
                    cell_list.append("-".join([cell_text, str(_index)]))
                    _index += 1
                elif cell_text == "占比" and _index == 2:
                    cell_list.append("-".join([cell_text, str(_index)]))
                    _index = 1
                else:
                    cell_list.append(cell_text)
            if not cell_list:
                continue
            group_list.append(cell_list)
            if "合计" in cell_list:
                result_dict[_title].append(group_list)
                group_list = list()

    for k, v in result_dict.items():
        group_list = list()
        for matrix in v:
            header = matrix[0]
            data = matrix[1:]
            df = pd.DataFrame(data, columns=header)
            for index, row in df.iterrows():
                key_value_pairs = row.to_dict()
                group_list.append(key_value_pairs)
        final_dict[k] = group_list
    
    return final_dict

def main(input_file, output_file):
    script_dir = os.path.dirname(os.path.abspath(__file__))
    input_file_path = os.path.join(script_dir + "/input_data/", input_file)
    output_file_path = os.path.join(script_dir+ "/output_data/", output_file)

    final_dict = load_docx2json(input_file_path)
    if os.path.exists(output_file_path):
        os.remove(output_file_path)
    with open(output_file_path, 'w', encoding='utf-8') as f:
        json.dump(final_dict, f, ensure_ascii=False, indent=4)


if __name__ == "__main__":
    input_file = input("请输入转换文件名称: ")
    print(f"输入文件名称为: {input_file}")
    output_file = input_file.split(".")[0] + ".json"
    print(f"输出文件名称为: {output_file}")
    try:
        main(input_file, output_file)
        time.sleep(2)
        print()
        print("...转换成功...OK!")
    except Exception as e:
        time.sleep(2)
        print()
        print("...转换失败...Failed!")
        print(e)
